#include "vkblob.h"
#include "vkutils.h"
#include "vkinfer/core/utils.h"
#include "vkcontext.h"
#include "vktransfer.h"

namespace vkinfer
{
	VkBlob::VkBlob(VkDevice device, VkPhysicalDevice gpu, const std::shared_ptr<VkTransfer>& transfer,
		BlobTransferDirection direction)
	{
		this->device = device;
		this->gpu = gpu;
		this->transfer = transfer;
		this->direction = direction;
	}

	VkBlob::VkBlob(VkDevice device, VkPhysicalDevice gpu, const std::shared_ptr<VkTransfer>& transfer, 
		const std::vector<uint32_t>& shape, BlobTransferDirection direction)
	{
		this->device = device;
		this->gpu = gpu;
		this->transfer = transfer;
		this->direction = direction;
		this->create(shape);
	}

	VkBlob::~VkBlob()
	{
		// free memory
		this->destory();
	}

	void VkBlob::copy_from(const std::vector<float>& data)
	{
		uint32_t count = shape_size(this->shape);
		if (count != data.size())
		{
			std::vector<uint32_t> shape;
			shape.push_back((uint32_t)data.size());
			resize(shape);
		}

		vkDeviceWaitIdle(device);
		void* mappedMemory = NULL;
		vkMapMemory(device, stagingMemory, 0, size, 0, &mappedMemory);
		memcpy(mappedMemory, data.data(), size);
		vkUnmapMemory(device, stagingMemory);
		transfer->copy(stagingBuffer, buffer, size);
	}

	void VkBlob::copy_from(const std::shared_ptr<Tensor<float>>& tensor)
	{
		this->copy_from(tensor->data, tensor->shape);
	}

	void VkBlob::copy_from(const float* data, const std::vector<uint32_t>& shape)
	{
		uint32_t count = shape_size(this->shape);
		uint32_t newcount = shape_size(shape);
		if (count != newcount)
		{
			resize(shape);
		}
		vkDeviceWaitIdle(device);
		void* mappedMemory = NULL;
		vkMapMemory(device, stagingMemory, 0, size, 0, &mappedMemory);
		memcpy(mappedMemory, data, size);
		vkUnmapMemory(device, stagingMemory);
		transfer->copy(stagingBuffer, buffer, size);
	}

	void VkBlob::copy_to(std::vector<float>& data)
	{
		vkDeviceWaitIdle(device);
		transfer->copy(buffer, stagingBuffer, size);

		void* mappedMemory = NULL;
		vkMapMemory(device, stagingMemory, 0, size, 0, &mappedMemory);
		uint32_t count = shape_size(this->shape);
		data.resize(count);
		memcpy(data.data(), mappedMemory, size);
		vkUnmapMemory(device, stagingMemory);
	}

	void VkBlob::resize(const std::vector<uint32_t>& shape)
	{
		if (shape == this->shape)
			return;

		this->destory();
		this->shape = shape;
		this->create(this->shape);
	}

	void VkBlob::create(const std::vector<uint32_t>& shape)
	{
		uint32_t count = shape_size(shape);
		this->size = count * sizeof(float);

		VkBufferUsageFlags buffer_usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT;
		VkBufferUsageFlags staging_usage = 0;
		if (direction == BlobTransferDirection::BLOB_TRANSFER_SRC)
		{
			buffer_usage |= VK_BUFFER_USAGE_TRANSFER_SRC_BIT;
			staging_usage |= VK_BUFFER_USAGE_TRANSFER_DST_BIT;
		}
		else if (direction == BlobTransferDirection::BLOB_TRANSFER_DST)
		{
			buffer_usage |= VK_BUFFER_USAGE_TRANSFER_DST_BIT;
			staging_usage |= VK_BUFFER_USAGE_TRANSFER_SRC_BIT;
		}

		create_buffer(gpu, device, size,
			buffer_usage,
			VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
			buffer, memory);

#if ALLOW_PARTIAL_FORWARD
		staging_usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT;
		create_buffer(gpu, device, size,
			staging_usage,
			VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
			stagingBuffer, stagingMemory);
#else
		if (direction != BlobTransferDirection::BLOB_TRANSFER_NONE)
		{
			create_buffer(gpu, device, size,
				staging_usage,
				VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
				stagingBuffer, stagingMemory);
		}
#endif
	}

	void VkBlob::destory()
	{
		if (stagingMemory != VK_NULL_HANDLE)
			vkFreeMemory(device, stagingMemory, nullptr);
		if (stagingBuffer != VK_NULL_HANDLE)
			vkDestroyBuffer(device, stagingBuffer, nullptr);
		if (memory != VK_NULL_HANDLE)
			vkFreeMemory(device, memory, nullptr);
		if (buffer != VK_NULL_HANDLE)
			vkDestroyBuffer(device, buffer, nullptr);
	}
}
